import torch.nn as nn
import torch.nn.functional as F
import torch as th
from torch.distributions import Categorical
from torch.distributions import MultivariateNormal


class SepMLPAgent(nn.Module):
    def __init__(self, input_shape, args, id):
        super(SepMLPAgent, self).__init__()
        self.args = args
        self.id = id
        self.hidden_sizes = getattr(args, "hidden_sizes", [256, 256])

        self.fc1 = nn.Linear(input_shape, self.hidden_sizes[0])
        self.fc2 = nn.Linear(self.hidden_sizes[0], self.hidden_sizes[1])
        self.fc3 = nn.Linear(self.hidden_sizes[1], args.n_actions)

        with th.no_grad():
            th.nn.init.kaiming_normal_(self.fc1.weight, mode='fan_in', nonlinearity='relu')
            th.nn.init.kaiming_normal_(self.fc2.weight, mode='fan_in', nonlinearity='relu')
            th.nn.init.uniform_(self.fc3.weight, a=-1e-3, b=1e-3)

    def forward(self, inputs, hidden_state=None):
        x = F.relu(self.fc1(inputs))
        x = F.relu(self.fc2(x))
        pi = self.fc3(x)

        return pi, hidden_state
        
    def init_hidden(self):
        return None
